#include <iostream>
#include <chrono>
#ifdef USE_NEON
#include <arm_neon.h> // 添加ARM NEON头文件
#endif

static void BGR2GRAY_NoNEON(const unsigned char *bgr, int width, int height, unsigned char *gray) {
    int coef_q8[3] = { 29, 150, 77 };
    const unsigned char *temp_s = bgr;
    unsigned char *temp_d = gray;

    for (int row = 0; row < height; row++) {
        for (int col = 0; col < width; col++) {
            temp_d[0] = (temp_s[0] * coef_q8[0] + temp_s[1] * coef_q8[1] + temp_s[2] * coef_q8[2]) >> 8;
            temp_d += 1;
            temp_s += 3;
        }
    }
}

#ifdef USE_NEON
static void BGR2GRAY_NEON(const unsigned char *bgr, int width, int height, unsigned char *gray) {
    int coef_q8[3] = { 29, 150, 77 };
    const unsigned char *temp_s = bgr;
    unsigned char *temp_d = gray;

    uint8x8_t vwr = vdup_n_u8(coef_q8[0]);
    uint8x8_t vwg = vdup_n_u8(coef_q8[1]);
    uint8x8_t vwb = vdup_n_u8(coef_q8[2]);

    for (int row = 0; row < height; row++) {
        int col = 0;
        for (; col < width - 8; col += 8) {
            uint8x8x3_t vsrc = vld3_u8(temp_s);
            uint16x8_t vsum = vmull_u8(vsrc.val[0], vwr);
            vsum = vmlal_u8(vsum, vsrc.val[1], vwg);
            vsum = vmlal_u8(vsum, vsrc.val[2], vwb);
            vst1_u8(temp_d, vshrn_n_u16(vsum, 8));
            temp_d += 8;
            temp_s += 24;
        }
        for (; col < width; col++) {
            temp_d[0] = (temp_s[0] * coef_q8[0] + temp_s[1] * coef_q8[1] + temp_s[2] * coef_q8[2]) >> 8;
            temp_d += 1;
            temp_s += 3;
        }
    }
}
#endif

int main() {
    const int width = 1920, height = 1080;
    unsigned char *bgr = new unsigned char[width * height * 3];
    unsigned char *gray = new unsigned char[width * height];

    // 初始化测试数据
    for (int i = 0; i < width * height * 3; i++) {
        bgr[i] = static_cast<unsigned char>(i % 256);
    }

    // 测试 NoNEON 性能
    auto start = std::chrono::high_resolution_clock::now();
    BGR2GRAY_NoNEON(bgr, width, height, gray);
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "BGR2GRAY_NoNEON Time: " 
              << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() 
              << " ms" << std::endl;

#ifdef USE_NEON
    // 测试 NEON 性能
    start = std::chrono::high_resolution_clock::now();
    BGR2GRAY_NEON(bgr, width, height, gray);
    end = std::chrono::high_resolution_clock::now();
    std::cout << "BGR2GRAY_NEON Time: " 
              << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() 
              << " ms" << std::endl;
#endif

    delete[] bgr;
    delete[] gray;

    return 0;
}
